package parser;

import lexer.Lexer;
import lexer.Token;
import lexer.TokenType;
import java.util.ArrayList;
import java.util.List;

public class Parser {
    private Lexer lexer;
    private Token currentToken;
    private List<Token> tokens;
    private int currentTokenIndex;

    public Parser(Lexer lexer) {
        this.lexer = lexer;
        this.tokens = new ArrayList<>();
        this.currentTokenIndex = 0;

        // 一次性读取所有token
        Token token;
        do {
            token = lexer.getNextToken();
            tokens.add(token);
        } while (token.getType() != TokenType.EOF);

        this.currentToken = tokens.get(0);
    }

    private void eat(TokenType type) {
        if (currentToken.getType() == type) {
            currentTokenIndex++;
            if (currentTokenIndex < tokens.size()) {
                currentToken = tokens.get(currentTokenIndex);
            } else {
                currentToken = new Token(TokenType.EOF, "", currentToken.getLine(), currentToken.getColumn());
            }
        } else {
            throw new RuntimeException("Syntax error: expected " + type +
                    ", got " + currentToken.getType() +
                    " at line " + currentToken.getLine() +
                    ", value: '" + currentToken.getValue() + "'");
        }
    }

    // 查看下一个token而不消费它
    private Token peekNextToken() {
        if (currentTokenIndex + 1 < tokens.size()) {
            return tokens.get(currentTokenIndex + 1);
        }
        return new Token(TokenType.EOF, "", currentToken.getLine(), currentToken.getColumn());
    }

    private ASTNode program() {
        ASTNode node = new ASTNode(NodeType.PROGRAM);
        node.addChild(statementList());
        return node;
    }

    private ASTNode statementList() {
        ASTNode node = new ASTNode(NodeType.STATEMENT_LIST);

        while (currentToken.getType() != TokenType.EOF) {
            node.addChild(statement());

            // 语句之间可能用分号分隔
            if (currentToken.getType() == TokenType.SEMICOLON) {
                eat(TokenType.SEMICOLON);
            }
        }

        return node;
    }

    private ASTNode statement() {
        // 如果是标识符，并且下一个token是赋值运算符，则是赋值语句
        if (currentToken.getType() == TokenType.IDENTIFIER &&
                peekNextToken().getType() == TokenType.ASSIGN) {
            return assignmentStatement();
        }

        // 其他语句类型
        switch (currentToken.getType()) {
            case IF:
                return ifStatement();
            case FOR:
                return forLoop();
            case WHILE:
                return whileLoop();
            case BREAK:
                return breakStatement();
            case CONTINUE:
                return continueStatement();
            default:
                return expressionStatement();
        }
    }

    private ASTNode assignmentStatement() {
        Token variable = currentToken;
        eat(TokenType.IDENTIFIER);
        eat(TokenType.ASSIGN);

        ASTNode node = new ASTNode(NodeType.ASSIGNMENT_STATEMENT, variable);
        node.addChild(expression());
        return node;
    }

    private ASTNode expressionStatement() {
        ASTNode expr = expression();

        // 表达式语句可能以分号结束，但分号不是表达式的一部分
        if (currentToken.getType() == TokenType.SEMICOLON) {
            eat(TokenType.SEMICOLON);
        }

        return expr;
    }

    private ASTNode ifStatement() {
        eat(TokenType.IF);
        ASTNode condition = expression();

        // 通达信公式中THEN可能是可选的
        if (currentToken.getType() == TokenType.THEN) {
            eat(TokenType.THEN);
        }

        ASTNode thenBranch = statement();

        ASTNode elseBranch = null;
        if (currentToken.getType() == TokenType.ELSE) {
            eat(TokenType.ELSE);
            elseBranch = statement();
        }

        if (currentToken.getType() == TokenType.ENDIF) {
            eat(TokenType.ENDIF);
        }

        ASTNode node = new ASTNode(NodeType.IF_STATEMENT);
        node.addChild(condition);
        node.addChild(thenBranch);
        if (elseBranch != null) {
            node.addChild(elseBranch);
        }
        return node;
    }

    private ASTNode forLoop() {
        eat(TokenType.FOR);

        Token variable = currentToken;
        eat(TokenType.IDENTIFIER);
        eat(TokenType.ASSIGN);

        ASTNode start = expression();
        eat(TokenType.TO);
        ASTNode end = expression();

        ASTNode step = null;
        if (currentToken.getType() == TokenType.STEP) {
            eat(TokenType.STEP);
            step = expression();
        }

        if (currentToken.getType() == TokenType.DO) {
            eat(TokenType.DO);
        }

        ASTNode body = statement();

        if (currentToken.getType() == TokenType.ENDFOR) {
            eat(TokenType.ENDFOR);
        }

        ASTNode node = new ASTNode(NodeType.FOR_LOOP, variable);
        node.addChild(start);
        node.addChild(end);
        if (step != null) {
            node.addChild(step);
        }
        node.addChild(body);
        return node;
    }

    private ASTNode whileLoop() {
        eat(TokenType.WHILE);
        ASTNode condition = expression();

        if (currentToken.getType() == TokenType.DO) {
            eat(TokenType.DO);
        }

        ASTNode body = statement();

        if (currentToken.getType() == TokenType.ENDWHILE) {
            eat(TokenType.ENDWHILE);
        }

        ASTNode node = new ASTNode(NodeType.WHILE_LOOP);
        node.addChild(condition);
        node.addChild(body);
        return node;
    }

    private ASTNode breakStatement() {
        Token token = currentToken;
        eat(TokenType.BREAK);

        // 分号是可选的
        if (currentToken.getType() == TokenType.SEMICOLON) {
            eat(TokenType.SEMICOLON);
        }

        return new ASTNode(NodeType.BREAK_STATEMENT, token);
    }

    private ASTNode continueStatement() {
        Token token = currentToken;
        eat(TokenType.CONTINUE);

        // 分号是可选的
        if (currentToken.getType() == TokenType.SEMICOLON) {
            eat(TokenType.SEMICOLON);
        }

        return new ASTNode(NodeType.CONTINUE_STATEMENT, token);
    }

    private ASTNode expression() {
        return logicalExpression();
    }

    private ASTNode logicalExpression() {
        ASTNode node = relationalExpression();

        while (currentToken.getType() == TokenType.AND ||
                currentToken.getType() == TokenType.OR) {
            Token token = currentToken;
            if (token.getType() == TokenType.AND) {
                eat(TokenType.AND);
            } else {
                eat(TokenType.OR);
            }

            ASTNode right = relationalExpression();
            ASTNode binaryOp = new ASTNode(NodeType.BINARY_OPERATION, token);
            binaryOp.addChild(node);
            binaryOp.addChild(right);
            node = binaryOp;
        }

        return node;
    }

    private ASTNode relationalExpression() {
        ASTNode node = additiveExpression();

        while (isRelationalOperator(currentToken.getType())) {
            Token token = currentToken;
            eat(token.getType());

            ASTNode right = additiveExpression();
            ASTNode binaryOp = new ASTNode(NodeType.BINARY_OPERATION, token);
            binaryOp.addChild(node);
            binaryOp.addChild(right);
            node = binaryOp;
        }

        return node;
    }

    private boolean isRelationalOperator(TokenType type) {
        return type == TokenType.EQ || type == TokenType.NE ||
                type == TokenType.LT || type == TokenType.LE ||
                type == TokenType.GT || type == TokenType.GE;
    }

    private ASTNode additiveExpression() {
        ASTNode node = multiplicativeExpression();

        while (currentToken.getType() == TokenType.PLUS ||
                currentToken.getType() == TokenType.MINUS) {
            Token token = currentToken;
            if (token.getType() == TokenType.PLUS) {
                eat(TokenType.PLUS);
            } else {
                eat(TokenType.MINUS);
            }

            ASTNode right = multiplicativeExpression();
            ASTNode binaryOp = new ASTNode(NodeType.BINARY_OPERATION, token);
            binaryOp.addChild(node);
            binaryOp.addChild(right);
            node = binaryOp;
        }

        return node;
    }

    private ASTNode multiplicativeExpression() {
        ASTNode node = unaryExpression();

        while (currentToken.getType() == TokenType.MULTIPLY ||
                currentToken.getType() == TokenType.DIVIDE ||
                currentToken.getType() == TokenType.MOD) {
            Token token = currentToken;
            if (token.getType() == TokenType.MULTIPLY) {
                eat(TokenType.MULTIPLY);
            } else if (token.getType() == TokenType.DIVIDE) {
                eat(TokenType.DIVIDE);
            } else {
                eat(TokenType.MOD);
            }

            ASTNode right = unaryExpression();
            ASTNode binaryOp = new ASTNode(NodeType.BINARY_OPERATION, token);
            binaryOp.addChild(node);
            binaryOp.addChild(right);
            node = binaryOp;
        }

        return node;
    }

    private ASTNode unaryExpression() {
        if (currentToken.getType() == TokenType.PLUS ||
                currentToken.getType() == TokenType.MINUS ||
                currentToken.getType() == TokenType.NOT) {
            Token token = currentToken;
            eat(token.getType());

            ASTNode node = new ASTNode(NodeType.UNARY_OPERATION, token);
            node.addChild(unaryExpression());
            return node;
        }

        return primary();
    }

    private ASTNode primary() {
        Token token = currentToken;

        switch (token.getType()) {
            case NUMBER:
                eat(TokenType.NUMBER);
                return new ASTNode(NodeType.CONSTANT, token);
            case STRING:
                eat(TokenType.STRING);
                return new ASTNode(NodeType.CONSTANT, token);
            case IDENTIFIER:
            case IF:    // 添加IF关键字处理，因为IF也可以作为函数名
            case AND:   // 添加AND关键字处理
            case OR:    // 添加OR关键字处理
            case NOT:   // 添加NOT关键字处理
                eat(token.getType());

                // 检查是否是函数调用
                if (currentToken.getType() == TokenType.LPAREN) {
                    return functionCall(token);
                } else {
                    return new ASTNode(NodeType.VARIABLE, token);
                }
            case LPAREN:
                eat(TokenType.LPAREN);
                ASTNode node = expression();
                eat(TokenType.RPAREN);
                return node;
            default:
                throw new RuntimeException("Unexpected token in expression: " +
                        token.getType() + " '" + token.getValue() +
                        "' at line " + token.getLine());
        }
    }

    private ASTNode functionCall(Token functionName) {
        eat(TokenType.LPAREN);

        ASTNode node = new ASTNode(NodeType.FUNCTION_CALL, functionName);

        if (currentToken.getType() != TokenType.RPAREN) {
            node.addChild(expression());

            while (currentToken.getType() == TokenType.COMMA) {
                eat(TokenType.COMMA);
                node.addChild(expression());
            }
        }

        eat(TokenType.RPAREN);
        return node;
    }

    public ASTNode parse() {
        return program();
    }
}